import os
from PIL import Image
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt

from model import get_crnn


class LayerActivation():
    """
        show cnn visualization result

    Args:
        cnn_layer(nn.Module) : given layer of a model
        channel(int) : out channels of that layer
        save_file(str) : output file that store shape & value of given 
                        `channel` of that `cnn_layer`
    """
    def __init__(self, cnn_layer, channel=0, save_file=None):
        self.hook = cnn_layer.register_forward_hook(self.hook_fn)
        self.feature = None
        self.channel = channel
        self.save_file = save_file
    
    def hook_fn(self, module, inp, outp):
        """
        foward hook function
        """
        # 第一个0是为了去掉Batch维度
        self.feature = outp[0, self.channel].detach().cpu().numpy()
        if self.save_file is not None:
            with open(self.save_file, 'w') as f:
                f.write('shape:{}, values:{}\n'.format(self.feature.shape, self.feature))

    def remove(self):
        """
        don't forget to call this func to remove the hook when you finished visualization
        """
        self.hook.remove()

def main():

    data_root = 'data/train_imgs/'
    # img_name = '20161013_005618_128_0193gray.jpg' # 0.74
    img_name = '20130202_215919_128_0193gray.jpg' # 0.19

    out_root = 'cnn_visuals/'
    if not os.path.exists(out_root):
        os.makedirs(out_root)
    save_file = os.path.join(out_root, img_name.replace('.jpg', '.txt'))
    img_path = os.path.join(data_root, img_name)
    model = get_crnn()
    # conv_feature = LayerActivation(model.cnn, channel=0, save_file=save_file)
    conv_feature = LayerActivation(model.cnn.layer2, channel=0)

    img = Image.open(img_path)
    plt.figure()
    plt.imshow(img)
    trans = torchvision.transforms.Compose([transforms.ToTensor()])
    img_tensor = trans(img).unsqueeze(dim=0)

    model(img_tensor)
    conv_feature.remove()
    print(conv_feature.feature.shape)

    plt.figure()
    plt.imshow(conv_feature.feature)
    plt.show()

if __name__ == "__main__":
    main()
